import os
from typing import List
from langchain.document_loaders import TextLoader
from langchain.text_splitter import CharacterTextSplitter
from langchain.embeddings import OllamaEmbeddings
from langchain.vectorstores import Chroma
from langchain.llms import Ollama
from langchain.chains import RetrievalQA
from langchain.prompts import PromptTemplate

class RAGSystem:
    """
    RAG（检索增强生成）系统
    
    RAG 是一种结合文档检索和语言模型的技术，主要包含以下步骤：
    1. 文档加载和分割：将长文档分割成小块，便于后续处理
    2. 向量化存储：将文档块转换为向量并存储，便于相似度检索
    3. 检索相关内容：根据用户问题检索相关文档片段
    4. 生成回答：将检索到的相关内容和用户问题一起发送给语言模型生成回答
    """

    def __init__(self, file_path: str, chunk_size: int = 256):
        """
        初始化 RAG 系统
        :param file_path: 知识库文本文件路径
        :param chunk_size: 文档分块大小，较小的块有助于精确检索，但可能缺少上下文
        """
        self.file_path = file_path
        self.chunk_size = chunk_size
        self.documents = None  # 原始文档
        self.chunks = None     # 分割后的文档块
        self.db = None        # 向量数据库
        self.qa_chain = None  # 问答链

    def load_and_split_document(self) -> None:
        """
        加载并分割文档
        
        这一步骤将文本文件加载为文档对象，并将其分割成更小的块。
        分割是必要的，因为：
        1. 向量数据库对每个块单独编码，便于精确检索
        2. 语言模型通常有输入长度限制
        3. 较小的文本块有助于找到最相关的内容
        """
        # 检查文件是否存在
        if not os.path.exists(self.file_path):
            raise FileNotFoundError(f"文件 {self.file_path} 不存在")

        # 加载文档
        loader = TextLoader(self.file_path)
        self.documents = loader.load()

        # 分割文档
        splitter = CharacterTextSplitter(
            chunk_size=self.chunk_size,
            chunk_overlap=20,  # 添加一些重叠，确保上下文的连续性
            separator="\n",
            length_function=len,
            is_separator_regex=False,
            keep_separator=True,
            strip_whitespace=True
        )
        self.chunks = splitter.split_documents(self.documents)

        # 打印文档信息
        self._print_document_info()

    def _print_document_info(self) -> None:
        """打印文档相关信息，用于调试和监控"""
        print("\n=== 文档信息 ===")
        print(f"总文档块数: {len(self.chunks)}")
        print("\n前两个文档块示例:")
        for i in range(min(2, len(self.chunks))):
            print(f"\n文档块 {i+1}:")
            print(f"内容: {self.chunks[i].page_content[:50]}...")
            print(f"元数据: {self.chunks[i].metadata}")

        # 计算预估存储大小
        vector_size = 768 * 4  # 每个向量的字节数 (768维向量，每维4字节)
        metadata_size = 100    # 估计的每个文档元数据大小
        total_size = len(self.chunks) * (vector_size + metadata_size)
        print(f"\n预估存储大小: {total_size/1024/1024:.2f} MB")

    def create_vector_store(self, persist_directory: str = "./vector_db") -> None:
        """
        创建向量数据库
        
        这一步骤将文档块转换为向量并存储，是实现语义检索的关键：
        1. 使用 Ollama 的 embedding 模型将文本转换为向量
        2. 将向量存储在 Chroma 数据库中
        3. 后续可以通过向量相似度搜索相关内容
        
        :param persist_directory: 向量数据库存储路径
        """
        print("\n=== 创建向量数据库 ===")
        
        # 配置 embedding 模型
        # nomic-embed-text 是专门用于文本嵌入的模型
        embedding = OllamaEmbeddings(
            model="nomic-embed-text",
            base_url="http://localhost:11434"
        )

        # 创建向量数据库
        self.db = Chroma.from_documents(
            documents=self.chunks,
            embedding=embedding,
            persist_directory=persist_directory
        )
        
        # 持久化存储
        self.db.persist()
        print(f"向量数据库已创建并保存至: {persist_directory}")

    def setup_qa_chain(self) -> None:
        """
        设置问答链
        
        问答链将检索和生成过程串联起来：
        1. 检索器根据问题找到相关文档块
        2. 提示模板将问题和相关文档组织成合适的格式
        3. 语言模型根据提供的上下文生成答案
        """
        print("\n=== 设置问答链 ===")
        
        # 初始化 Ollama LLM
        llm = Ollama(
            model="llama3.3:70b",  # 或其他你已安装的模型
            base_url="http://localhost:11434"
        )

        # 创建提示模板
        # 这个模板告诉语言模型如何使用检索到的上下文来回答问题
        prompt_template = """请基于以下已知信息，简洁专业地回答问题。
如果无法从中得到答案，请说"抱歉，根据已知信息无法回答该问题"。

已知信息：
{context}

问题：{question}
回答："""

        PROMPT = PromptTemplate(
            template=prompt_template,
            input_variables=["context", "question"]
        )

        # 创建检索问答链
        # k=3 表示每次检索3个最相关的文档块
        retriever = self.db.as_retriever(search_kwargs={"k": 3})
        self.qa_chain = RetrievalQA.from_chain_type(
            llm=llm,
            chain_type="stuff",  # "stuff" 策略将所有检索到的文档合并后一次性发送给LLM
            retriever=retriever,
            return_source_documents=True,  # 返回源文档便于追踪答案来源
            chain_type_kwargs={"prompt": PROMPT}
        )
        
        print("问答链设置完成")

    def ask_question(self, question: str) -> None:
        """
        提问并获取答案
        
        整个RAG过程的最后一步：
        1. 问题被发送给检索器找到相关文档
        2. 相关文档与问题一起通过提示模板格式化
        3. 语言模型生成最终答案
        
        :param question: 用户问题
        """
        print(f"\n=== 问答测试 ===")
        print(f"问题: {question}")
        
        # 执行问答链
        response = self.qa_chain({"query": question})
        print(f"\n答案: {response['result']}")
        
        # 打印相关文档来源，便于验证答案
        print("\n相关文档片段:")
        for i, doc in enumerate(response['source_documents'], 1):
            print(f"\n文档 {i}:")
            print(f"内容: {doc.page_content[:100]}...")

def main():
    """使用示例"""
    # 初始化RAG系统，指定知识库文件
    rag = RAGSystem("li.txt")
    
    try:
        # 1. 加载和分割文档
        rag.load_and_split_document()
        
        # 2. 创建向量数据库
        rag.create_vector_store()
        
        # 3. 设置问答链
        rag.setup_qa_chain()
        
        # 4. 提问测试
        rag.ask_question("藜怎么防治虫害？")
        
    except Exception as e:
        print(f"发生错误: {str(e)}")

if __name__ == "__main__":
    main()
